import argparse
import numpy as np
import torch
from collections import OrderedDict
from mmcv import Config
from models.backbones.repvgg import RepVGG
from pathlib import Path

from mmdet.models import build_backbone, build_detector
from mmhead import models  # noqa: F401,F403

optional_groupwise_layers = [2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26]
g2_map = {l: 2 for l in optional_groupwise_layers}
g4_map = {l: 4 for l in optional_groupwise_layers}


def create_RepVGG_A0(deploy=False, out_stage_names=list()):
    return RepVGG(
        num_blocks=[2, 4, 14, 1],
        width_multiplier=[0.75, 0.75, 0.75, 2.5],
        out_stage_names=out_stage_names,
        override_groups_map=None,
        deploy=deploy)


def create_RepVGG_A1(deploy=False, out_stage_names=list()):
    return RepVGG(
        num_blocks=[2, 4, 14, 1],
        width_multiplier=[1, 1, 1, 2.5],
        out_stage_names=out_stage_names,
        override_groups_map=None,
        deploy=deploy)


def create_RepVGG_A2(deploy=False, out_stage_names=list()):
    return RepVGG(
        num_blocks=[2, 4, 14, 1],
        width_multiplier=[1.5, 1.5, 1.5, 2.75],
        out_stage_names=out_stage_names,
        override_groups_map=None,
        deploy=deploy)


def create_RepVGG_B0(deploy=False, out_stage_names=list()):
    return RepVGG(
        num_blocks=[4, 6, 16, 1],
        width_multiplier=[1, 1, 1, 2.5],
        out_stage_names=out_stage_names,
        override_groups_map=None,
        deploy=deploy)


def create_RepVGG_B1(deploy=False, out_stage_names=list()):
    return RepVGG(
        num_blocks=[4, 6, 16, 1],
        width_multiplier=[2, 2, 2, 4],
        out_stage_names=out_stage_names,
        override_groups_map=None,
        deploy=deploy)


def create_RepVGG_B1g2(deploy=False, out_stage_names=list()):
    return RepVGG(
        num_blocks=[4, 6, 16, 1],
        width_multiplier=[2, 2, 2, 4],
        out_stage_names=out_stage_names,
        override_groups_map=g2_map,
        deploy=deploy)


def create_RepVGG_B1g4(deploy=False, out_stage_names=list()):
    return RepVGG(
        num_blocks=[4, 6, 16, 1],
        width_multiplier=[2, 2, 2, 4],
        out_stage_names=out_stage_names,
        override_groups_map=g4_map,
        deploy=deploy)


def create_RepVGG_B2(deploy=False, out_stage_names=list()):
    return RepVGG(
        num_blocks=[4, 6, 16, 1],
        width_multiplier=[2.5, 2.5, 2.5, 5],
        out_stage_names=out_stage_names,
        override_groups_map=None,
        deploy=deploy)


def create_RepVGG_B2g2(deploy=False, out_stage_names=list()):
    return RepVGG(
        num_blocks=[4, 6, 16, 1],
        width_multiplier=[2.5, 2.5, 2.5, 5],
        out_stage_names=out_stage_names,
        override_groups_map=g2_map,
        deploy=deploy)


def create_RepVGG_B2g4(deploy=False, out_stage_names=list()):
    return RepVGG(
        num_blocks=[4, 6, 16, 1],
        width_multiplier=[2.5, 2.5, 2.5, 5],
        out_stage_names=out_stage_names,
        override_groups_map=g4_map,
        deploy=deploy)


def create_RepVGG_B3(deploy=False, out_stage_names=list()):
    return RepVGG(
        num_blocks=[4, 6, 16, 1],
        width_multiplier=[3, 3, 3, 5],
        out_stage_names=out_stage_names,
        override_groups_map=None,
        deploy=deploy)


def create_RepVGG_B3g2(deploy=False, out_stage_names=list()):
    return RepVGG(
        num_blocks=[4, 6, 16, 1],
        width_multiplier=[3, 3, 3, 5],
        out_stage_names=out_stage_names,
        override_groups_map=g2_map,
        deploy=deploy)


def create_RepVGG_B3g4(deploy=False, out_stage_names=list()):
    return RepVGG(
        num_blocks=[4, 6, 16, 1],
        width_multiplier=[3, 3, 3, 5],
        out_stage_names=out_stage_names,
        override_groups_map=g4_map,
        deploy=deploy)


func_dict = {
    'RepVGG-A0': create_RepVGG_A0,
    'RepVGG-A1': create_RepVGG_A1,
    'RepVGG-A2': create_RepVGG_A2,
    'RepVGG-B0': create_RepVGG_B0,
    'RepVGG-B1': create_RepVGG_B1,
    'RepVGG-B1g2': create_RepVGG_B1g2,
    'RepVGG-B1g4': create_RepVGG_B1g4,
    'RepVGG-B2': create_RepVGG_B2,
    'RepVGG-B2g2': create_RepVGG_B2g2,
    'RepVGG-B2g4': create_RepVGG_B2g4,
    'RepVGG-B3': create_RepVGG_B3,
    'RepVGG-B3g2': create_RepVGG_B3g2,
    'RepVGG-B3g4': create_RepVGG_B3g4,
}


def get_RepVGG_func_by_name(name):
    return func_dict[name]


#   Use like this:
#   train_model = create_RepVGG_A0(deploy=False)
#   train train_model
#   deploy_model = repvgg_convert(train_model, create_RepVGG_A0, save_path='repvgg_deploy.pth')
def repvgg_model_convert(model: torch.nn.Module,
                         build_func,
                         out_stage_names=list(),
                         save_path=None):
    converted_weights = {}
    for name, module in model.named_modules():
        if hasattr(module, 'repvgg_convert'):
            kernel, bias = module.repvgg_convert()
            converted_weights[name + '.rbr_reparam.weight'] = kernel
            converted_weights[name + '.rbr_reparam.bias'] = bias
        elif isinstance(module, torch.nn.Linear):
            converted_weights[
                name + '.weight'] = module.weight.detach().cpu().numpy()
            converted_weights[name +
                              '.bias'] = module.bias.detach().cpu().numpy()
        else:
            print(name, type(module))
    del model

    deploy_model = build_func(deploy=True, out_stage_names=out_stage_names)
    for name, param in deploy_model.named_parameters():
        print('deploy param: ', name, param.size(),
              np.mean(converted_weights[name]))
        param.data = torch.from_numpy(converted_weights[name]).float()

    if save_path is not None:
        torch.save(deploy_model.state_dict(), save_path)

    return deploy_model


def argument_parser():
    parser = argparse.ArgumentParser(description="convert repvgg")
    parser.add_argument(
        "--config-file",
        default="",
        metavar="FILE",
        help="path to config file")
    parser.add_argument(
        "--checkpoint",
        default="",
        metavar="FILE",
        help="path to checkpoint file")
    parser.add_argument(
        "--arch",
        default="RepVGG-A1",
        help="repvgg architecture name, default RepVGG-A1")
    args = parser.parse_args()
    return args


def main():
    args = argument_parser()
    config_file_path = args.config_file
    checkpoint_path = args.checkpoint
    arch_name = args.arch

    cfg = Config.fromfile(config_file_path)
    checkpoint = torch.load(checkpoint_path)

    out_stage_names = cfg.model.backbone.out_stage_names
    repvgg_build_func = get_RepVGG_func_by_name(arch_name)
    backbone_train_model = repvgg_build_func(
        deploy=False, out_stage_names=out_stage_names)
    backbone_train_model.gap = None
    backbone_train_model.linear = None

    backbone_state_dict = OrderedDict()
    for k, v in checkpoint["state_dict"].items():
        if "backbone" in k:
            backbone_state_dict[k.replace("backbone.", "")] = v
    backbone_train_model.load_state_dict(backbone_state_dict, strict=True)
    backbone_deploy_model = repvgg_model_convert(backbone_train_model,
                                                 repvgg_build_func,
                                                 out_stage_names)
    backbone_deploy_model = backbone_deploy_model.state_dict(
        prefix="backbone.")

    cfg.model.backbone.deploy = True
    cfg.model.pretrained = None
    detector = build_detector(cfg.model)
    results = detector.load_state_dict(backbone_deploy_model, strict=False)
    assert len(results.unexpected_keys) == 0
    results = detector.load_state_dict(checkpoint["state_dict"], strict=False)

    checkpoint.pop("state_dict")
    checkpoint["state_dict"] = detector.state_dict()

    checkpoint_path = Path(checkpoint_path)
    save_name = checkpoint_path.stem + "_deploy" + checkpoint_path.suffix
    save_path = str(checkpoint_path.parent / save_name)
    torch.save(checkpoint, save_path)


if __name__ == "__main__":
    main()
